package cn.ctyun.xstore.s3.controller;

import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
import com.amazonaws.services.securitytoken.model.AssumeRoleRequest;
import com.amazonaws.services.securitytoken.model.AssumeRoleResult;
import com.amazonaws.services.securitytoken.model.Credentials;

import cn.ctyun.xstore.s3.Conf;
import cn.ctyun.xstore.s3.api.object.GeneratePresignedUrl;
import cn.ctyun.xstore.s3.client.AwsS3ClientBuilder;
import cn.ctyun.xstore.s3.client.AwsSTSClientBuilder;

import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

@RestController
@RequestMapping(value = "/sts")
public class StsDemoController {

	private static final String ARN = "";

	private static final String POLICY = "{\"Version\":\"2012-10-17\"," + "\"Statement\":" + "{\"Effect\":\"Allow\","
			+ "\"Action\":[\"s3:*\"]," // 允许进行 S3 的所有操作。如果仅需要上传，这里可以设置为 PutObject
			+ "\"Resource\":[\"arn:aws:s3:::" + Conf.DEFAULT_BUCKET + "\",\"arn:aws:s3:::" + Conf.DEFAULT_BUCKET + "/*\"]"// 允许操作默认桶中的所有文件，可以修改此处来保证操作的文件
			+ "}}";

	/**
	 * 取得一个临时的AK/SK给页面使用
	 *
	 * @return
	 * @throws IOException
	 */
	@RequestMapping(value = "/assumeRole")
	public Map<String, String> assumeRole() throws IOException {
		Map<String, String> map = new HashMap<>();
		try {
			AWSSecurityTokenService stsClient = AwsSTSClientBuilder.buildSTSClient();
			AssumeRoleRequest assumeRoleRequest = new AssumeRoleRequest();
			assumeRoleRequest.setRoleArn(ARN);
			assumeRoleRequest.setPolicy(POLICY);
			assumeRoleRequest.setRoleSessionName("S3AccessSession");
			AssumeRoleResult assumeRoleRes = stsClient.assumeRole(assumeRoleRequest);
			Credentials stsCredentials = assumeRoleRes.getCredentials();
			map.put("ak", stsCredentials.getAccessKeyId());
			map.put("sk", stsCredentials.getSecretAccessKey());
			map.put("tk", stsCredentials.getSessionToken());
			map.put("endpoint", Conf.END_POINT);
			map.put("bucket", Conf.DEFAULT_BUCKET);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return map;
	}

	/**
	 * ִ执行流程： * 1、根据正式ak、sk生成临时ak、sk、token * 2、根据临时ak、sk、token生成文件上传预签名地址
	 *
	 * @return
	 * @throws IOException
	 */
	@RequestMapping(value = "/presignedPutUrl")
	public Map<String, String> generatePresignedPutUrl(String key) throws IOException {
		String url = "";
		Map<String, String> map = new HashMap<>();
		try {
			// 根据正式ak、sk生成临时ak、sk、token
			Map<String, String> sessionCredentials = assumeRole();
			// 根据临时ak、sk、token生成文件上传预签名地址
			AmazonS3 s3Client = AwsS3ClientBuilder.buildSessionCredentialsAmazonS3Client(
					sessionCredentials.get("endpoint"), sessionCredentials.get("ak"), sessionCredentials.get("sk"),
					sessionCredentials.get("tk"));
			url = GeneratePresignedUrl.generatePresignedPutUrl(key, sessionCredentials.get("bucket"), s3Client);
			map.put("url", url);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return map;
	}
}
